from motors.feetech import FeetechMotorsBus
from motors.configs import FeetechMotorsBusConfig
from motors.feetech import TorqueMode
from vr_monitor import VRMonitor
import time
import threading
import asyncio
import numpy as np
import math
import pygame

# 配置串口参数和电机信息
config1 = FeetechMotorsBusConfig(
    port="/dev/ttyACM0",  # 串口端口
    motors={
            # name: (index, model)
            "shoulder_pan": [1, "sts3215"],
            "shoulder_lift": [2, "sts3215"],
            "elbow_flex": [3, "sts3215"],
            "wrist_flex": [4, "sts3215"],
            "wrist_roll": [5, "sts3215"],
            "gripper": [6, "sts3215"],
            }
)

config2 = FeetechMotorsBusConfig(
    port="/dev/ttyACM1",  # 串口端口
    motors={
            # name: (index, model)
            "shoulder_pan": [1, "sts3215"],
            "shoulder_lift": [2, "sts3215"],
            "elbow_flex": [3, "sts3215"],
            "wrist_flex": [4, "sts3215"],
            "wrist_roll": [5, "sts3215"],
            "gripper": [6, "sts3215"],
            }
)



def get_mapped_joints(left,right):
    all_positions_left = left.read("Present_Position")
    all_positions_right = right.read("Present_Position")
    mapped_joints = np.zeros(16)
    mapped_joints[0] = 0  # Base X position
    mapped_joints[1] = 0  # Base rotation
    
    # First arm: [3,6,9,11,13] → [2,3,4,5,6]
    mapped_joints[2] = all_positions_left[0]
    mapped_joints[3] = all_positions_left[1]
    mapped_joints[4] = all_positions_left[2]
    mapped_joints[5] = all_positions_left[3]
    mapped_joints[6] = all_positions_left[4]
    
    # Second arm: [4,7,10,12,14] → [7,8,9,10,11]
    mapped_joints[7] = all_positions_right[0]
    mapped_joints[8] = all_positions_right[1]
    mapped_joints[9] = all_positions_right[2]
    mapped_joints[10] = all_positions_right[3]
    mapped_joints[11] = all_positions_right[4]
    mapped_joints[12] = all_positions_left[5]
    mapped_joints[13] = all_positions_right[5]
    
    return mapped_joints

def inverse_kinematics(x, y, l1=0.1159, l2=0.1350):

    # Calculate joint2 and joint3 offsets in theta1 and theta2
    theta1_offset = -math.atan2(0.028, 0.11257)  # theta1 offset when joint2=0
    theta2_offset = -math.atan2(0.0052, 0.1349) + theta1_offset  # theta2 offset when joint3=0
    
    # Calculate distance from origin to target point
    r = math.sqrt(x**2 + y**2)
    r_max = l1 + l2  # Maximum reachable distance
    
    # If target point is beyond maximum workspace, scale it to the boundary
    if r > r_max:
        scale_factor = r_max / r
        x *= scale_factor
        y *= scale_factor
        r = r_max
    
    # If target point is less than minimum workspace (|l1-l2|), scale it
    r_min = abs(l1 - l2)
    if r < r_min and r > 0:
        scale_factor = r_min / r
        x *= scale_factor
        y *= scale_factor
        r = r_min
    
    # Use law of cosines to calculate theta2
    cos_theta2 = -(r**2 - l1**2 - l2**2) / (2 * l1 * l2)
    
    # Calculate theta2 (elbow angle)
    theta2 = math.pi - math.acos(cos_theta2)
    
    # Calculate theta1 (shoulder angle)
    beta = math.atan2(y, x)
    gamma = math.atan2(l2 * math.sin(theta2), l1 + l2 * math.cos(theta2))
    theta1 = beta + gamma
    
    # Convert theta1 and theta2 to joint2 and joint3 angles
    joint2 = theta1 - theta1_offset
    joint3 = theta2 - theta2_offset
    
    # Ensure angles are within URDF limits
    joint2 = max(-0.1, min(3.45, joint2))
    joint3 = max(-0.2, min(math.pi, joint3))
    
    return joint2, joint3



def main():
    pygame.init()
    # 初始化控制窗口（仅保留控制交互功能）
    screen_width, screen_height = 600, 750
    screen = pygame.display.set_mode((screen_width, screen_height))
    pygame.display.set_caption("双机械臂控制窗口")
    font = pygame.font.SysFont(None, 24)
    

    # 初始化VR监控器
    vr_monitor = VRMonitor()
    # 启动独立线程运行VR监控（异步获取数据）
    vr_thread = threading.Thread(target=lambda: asyncio.run(vr_monitor.start_monitoring()),daemon=True)
    vr_thread.start()
    #    等待监控器初始化
    time.sleep(2)
    # 获取最新VR数据（左控制器、右控制器、头显）
    # 初始化电机总线并建立串口连接

    motors_bus_left = FeetechMotorsBus(config1)
    motors_bus_right = FeetechMotorsBus(config2)
    motors_bus_left.connect()  # 打开串口，建立通信
    motors_bus_right.connect()  # 打开串口，建立通信
    motors_bus_left.write("Torque_Enable", TorqueMode.ENABLED.value)
    motors_bus_right.write("Torque_Enable", TorqueMode.ENABLED.value)


    # VR映射缩放因子（根据真实机械臂工作空间调整）
    vr_scale_y = 1.0
    vr_scale_x = 0.5
    vr_scale_z = -0.6
    
    # 初始化关节目标位置
    action = np.zeros(16)
    target_joints = np.zeros(16)
    target_joints[3] = 2048
    target_joints[4] = 2048
    target_joints[8] = 2048
    target_joints[9] = 2048
    
    # 末端执行器初始位置
    initial_ee_pos_arm1 = np.array([-0.2, 0])  # 左机械臂
    initial_ee_pos_arm2 = np.array([0.2, 0])  # 右机械臂
    ee_pos_arm1 = initial_ee_pos_arm1.copy()
    ee_pos_arm2 = initial_ee_pos_arm2.copy()
    
    # 末端执行器姿态调整参数
    initial_pitch_1 = 0.0
    initial_pitch_2 = 0.0
    pitch_1 = initial_pitch_1
    pitch_2 = initial_pitch_2
    pitch_step = 0.02
    
    # 比例控制器增益（根据真实机械臂特性调整）
    p_gain = np.ones(16)
    p_gain[0] = 0    # 基座前后（若有）
    p_gain[1] = 0    # 基座旋转（若有）
    p_gain[2:7] = 1.0  # 左机械臂关节
    p_gain[7:12] = 1.0 # 右机械臂关节
    p_gain[12:14] = 0.1 # 夹爪关节
    p_gain[14:16] = 2.0 # 预留关节
    
    # 初始化关节位置
    current_joints = get_mapped_joints(motors_bus_left, motors_bus_right)
    target_joints = np.zeros_like(current_joints)
    
    # # 初始逆运动学计算
    # try:
    #     target_joints[3], target_joints[4] = inverse_kinematics(ee_pos_arm1[0], ee_pos_arm1[1])
    #     target_joints[8], target_joints[9] = inverse_kinematics(ee_pos_arm2[0], ee_pos_arm2[1])
    # except Exception as e:
    #     print(f"初始逆运动学计算错误: {e}")
    
    # 热身步骤计数器
    step_counter = 0
    warmup_steps = 50
    
    # 夹爪状态跟踪
    last_gripper_state_left = False
    last_gripper_state_right = False
    
    # 自动重置功能
    last_vr_activity_time = time.time()
    auto_reset_timeout = 5.0
    
    def reset_positions():
        """重置所有位置到初始状态"""
        nonlocal ee_pos_arm1, ee_pos_arm2, pitch_1, pitch_2, target_joints
        nonlocal last_gripper_state_left, last_gripper_state_right
        
        ee_pos_arm1 = initial_ee_pos_arm1.copy()
        ee_pos_arm2 = initial_ee_pos_arm2.copy()
        pitch_1 = initial_pitch_1
        pitch_2 = initial_pitch_2
        target_joints = np.zeros_like(target_joints)
        last_gripper_state_left = False
        last_gripper_state_right = False
        
        try:
            compensated_y1 = ee_pos_arm1[1] - 0.108 * math.sin(pitch_1)
            target_joints[3], target_joints[4] = inverse_kinematics(ee_pos_arm1[0], compensated_y1)
            target_joints[5] = target_joints[3] - target_joints[4] + pitch_1
            
            compensated_y2 = ee_pos_arm2[1] - 0.108 * math.sin(pitch_2)
            target_joints[8], target_joints[9] = inverse_kinematics(ee_pos_arm2[0], compensated_y2)
            target_joints[10] = target_joints[8] - target_joints[9] + pitch_2
        except Exception as e:
            print(f"重置时逆运动学错误: {e}")
        
        print("所有位置已重置到初始值")
    
    # 主控制循环
    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                # 断开机械臂连接
                motors_bus_left.disconnect()
                motors_bus_right.disconnect()
                # 停止VR监控
                if vr_monitor.is_running:
                    asyncio.run(vr_monitor.stop_monitoring())
                return
        
        # 获取VR控制器数据
        dual_goals = vr_monitor.get_latest_goal_nowait()
        # print(dual_goals)
        # print( )
        left_goal = dual_goals.get("left") if dual_goals else None
        # print( )
        # print( )
        # print(left_goal)
        right_goal = dual_goals.get("right") if dual_goals else None
        # print( )
        # print( )
        # print(right_goal)
        # print( )
        # print( )      
        
        # VR活动检测与自动重置
        current_time = time.time()
        vr_activity_detected = False
        kid = 11.11
        
        # 热身完成后处理VR输入
        if step_counter >= warmup_steps:
            # 处理左手控制器
            if left_goal is not None and left_goal.target_position is not None:
                vr_activity_detected = True
                pos = left_goal.target_position
                # x_vr = (pos[0] + 0.1) * vr_scale_x
                # y_vr = (pos[1] - 0.96) * vr_scale_y
                # z_vr = (pos[2] + 0.4) * vr_scale_z
                angle_rad = math.atan2(pos[2]+0.2, pos[0])

                # 计算旋转角度（左机械臂）
                target_joints[2] = angle_rad * kid +2048
                
                # 手腕弯曲角度（左机械臂）
                if left_goal.wrist_flex_deg is not None:
                    pitch_1 = (left_goal.wrist_flex_deg + 60) * 0.02
                
                # 手腕旋转角度（左机械臂）
                if left_goal.wrist_roll_deg is not None:
                    target_joints[6] = (-(left_goal.wrist_roll_deg - 90) * 0.02)* kid +2048
                
                # 左夹爪控制
                if left_goal.metadata.get('trigger', 0) > 0.5:
                    target_joints[12] = 3048  # 打开
                else:
                    target_joints[12] = 2048  # 闭合
                
            
            # 处理右手控制器
            if right_goal is not None and right_goal.target_position is not None:
                vr_activity_detected = True
                pos = right_goal.target_position
                
                # x_vr = (pos[0] - 0.1) * vr_scale_x
                # y_vr = (pos[1] - 0.96) * vr_scale_y
                # z_vr = (pos[2] + 0.4) * vr_scale_z
                angle_rad = math.atan2(pos[2]-0.2, pos[0])
                # 计算旋转角度（左机械臂）
                target_joints[7] = angle_rad * kid +2048
                
                # 手腕弯曲角度（右机械臂）
                if right_goal.wrist_flex_deg is not None:
                    pitch_2 = (right_goal.wrist_flex_deg + 60) * 0.02
                
                # 手腕旋转角度（右机械臂）
                if right_goal.wrist_roll_deg is not None:
                    target_joints[11] = (-(right_goal.wrist_roll_deg - 90) * 0.02)* kid +2048
                
                # 右夹爪控制
                if right_goal.metadata.get('trigger', 0) > 0.5:
                    target_joints[13] = 3048  # 打开
                else:
                    target_joints[13] = 2048  # 闭合
            
            # 计算逆运动学
            try:
                # 左机械臂
                compensated_y1 = ee_pos_arm1[1] + 0.108 * math.sin(pitch_1)
                target_joints[3], target_joints[4] = inverse_kinematics(ee_pos_arm1[0], compensated_y1)
                target_joints[5] = target_joints[3] - target_joints[4] + pitch_1
                
                # 右机械臂
                compensated_y2 = ee_pos_arm2[1] + 0.108 * math.sin(pitch_2)
                target_joints[8], target_joints[9] = inverse_kinematics(ee_pos_arm2[0], compensated_y2)
                target_joints[10] = target_joints[8] - target_joints[9] + pitch_2
            except Exception as e:
                print(f"逆运动学计算错误: {e}")
        
        # 读取当前关节位置
        current_joints = get_mapped_joints(motors_bus_left, motors_bus_right)
        
        # 比例控制计算
        action = p_gain * (target_joints - current_joints) 

        # 发送控制指令到真实机械臂
        # 提取左机械臂关节指令（索引2-6, 12）
        left_joints_cmd = [
            target_joints[2], target_joints[3]* kid +2048
                , target_joints[4]* kid +2048
                ,
            target_joints[5]* kid +2048
                , target_joints[6], target_joints[12]
        ]
        print("left")
        left_cmd = [int(x) for x in left_joints_cmd]
        print(left_cmd)

        motors_bus_left.write("Goal_Position", left_cmd)
        
        # 提取右机械臂关节指令（索引7-11, 13）
        right_joints_cmd = [
            target_joints[7], target_joints[8]* kid +2048
                , target_joints[9]* kid +2048
                ,
            target_joints[10]* kid +2048
                , target_joints[11], target_joints[13]
        ]
        print("right")
        right_cmd = [int(x) for x in right_joints_cmd]
        print(right_cmd)

        motors_bus_right.write("Goal_Position", right_cmd)
        
        # 更新活动时间
        if vr_activity_detected:
            last_vr_activity_time = current_time
        else:
            # 超时自动重置
            if current_time - last_vr_activity_time > auto_reset_timeout:
                reset_positions()
                last_vr_activity_time = current_time
        
        # 更新步骤计数器
        step_counter += 1
        
        # 刷新控制窗口
        screen.fill((255, 255, 255))
        status_text = font.render(f"当前关节位置: {current_joints[:6]}", True, (0, 0, 0))
        screen.blit(status_text, (10, 10))
        pygame.display.flip()
        
        # 控制循环频率（根据机械臂响应速度调整）
        time.sleep(0.5)

if __name__ == "__main__":
    main()